package llmserviceopenai.service

// RerankService.kt
import io.ktor.client.*
import io.ktor.client.call.body
import io.ktor.client.engine.cio.CIO
import io.ktor.client.plugins.ClientRequestException
import io.ktor.client.plugins.HttpTimeout
import io.ktor.client.plugins.ServerResponseException
import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
import io.ktor.client.request.header
import io.ktor.client.request.post
import io.ktor.client.request.setBody
import io.ktor.client.statement.HttpResponse
import io.ktor.client.statement.bodyAsText
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.contentType
import io.ktor.serialization.kotlinx.json.json
import kotlinx.serialization.json.Json
import llmserviceopenai.model.RerankRequest
import llmserviceopenai.model.RerankResponse
import llmserviceopenai.model.RerankResultItem
import llmserviceopenai.model.VLLMRerankResponse
import org.slf4j.LoggerFactory

class RerankService {

    private val logger = LoggerFactory.getLogger(RerankService::class.java)

    // 使用统一 JSON 配置（与你的 EmbeddingService 保持一致）
    private val json = Json {
        ignoreUnknownKeys = true
        explicitNulls = false
    }

    private val client = HttpClient(CIO) {
        install(ContentNegotiation) {
            json(json)
        }

        install(HttpTimeout) {
            requestTimeoutMillis = 600_00
            connectTimeoutMillis = 30_000
            socketTimeoutMillis = 600_00
        }
    }

    // vLLM Reranker 地址（注意：不是 /v1/rerank，而是 /rerank）
    private val backendUrl = "http://10.10.1.166:10011/rerank"
    private val apiKey = "EMPTY"  // 如果 vLLM 没设 key，可用 EMPTY
    private val modelName = "Qwen3-Reranker-4B"

    /**
     * 主接口：接收 query 和文档列表，返回按相关性排序的结果
     */
    suspend fun rerankHandler(request: RerankRequest): RerankResponse {
        logger.info("Forwarding rerank request to vLLM: query='${request.query}', documentCount=${request.documents.size}, topN=${request.topN}")

        if (request.documents.isEmpty()) {
            logger.warn("Empty documents list provided.")
            return RerankResponse(emptyList())
        }

        val startTime = System.currentTimeMillis()

        try {

            val response: HttpResponse = client.post(backendUrl) {
                contentType(ContentType.Application.Json)
                setBody(request)
                header(HttpHeaders.Authorization, "Bearer $apiKey")
            }

            val vllmResponse: VLLMRerankResponse = response.body()
            val duration = System.currentTimeMillis() - startTime

            // 转换为对外暴露的格式，并补全文本内容
            val resultItems = vllmResponse.results
                .sortedByDescending { it.relevance_score }
                .map { item ->
                    RerankResultItem(
                        index = item.index,
                        score = item.relevance_score,
                        text = request.documents.getOrNull(item.index) ?: "[missing]"
                    )
                }

            val finalResponse = RerankResponse(resultItems)

            logger.info("Reranking completed in $duration ms, top result score: ${resultItems.firstOrNull()?.score}")
            logger.debug("Rerank results: $finalResponse")

            return finalResponse

        } catch (e: ClientRequestException) {
            val status = e.response.status
            val errorBody = e.response.bodyAsText()
            logger.error("Client error during rerank ($status): $errorBody", e)
            throw e
        } catch (e: ServerResponseException) {
            val status = e.response.status
            val errorBody = e.response.bodyAsText()
            logger.error("Server error during rerank ($status): $errorBody", e)
            throw e
        } catch (e: Exception) {
            logger.error("Unexpected error while calling vLLM reranker", e)
            throw e
        }
    }

    /**
     * 关闭客户端资源
     */
    fun close() {
        client.close()
    }
}